import numpy as np

from Utils.inference import get_max_preds

def calc_dists(preds, target, normalize):
    '''
    计算两个张量的高斯核距离
    preds:shape=(batch_size,num_points,2)
    target:shape=(batch_size,num_points,2)
    normalize:shape=(1,2)
    '''
    preds = preds.astype(np.float32)
    target = target.astype(np.float32)
    dists = np.zeros((preds.shape[1], preds.shape[0]))
    for n in range(preds.shape[0]):
        for c in range(preds.shape[1]):
            #如果第n张图片的第c个点
            if target[n, c, 0] > 1 and target[n, c, 1] > 1:
                normed_preds = preds[n, c, :] / normalize[n]
                normed_targets = target[n, c, :] / normalize[n]
                dists[c, n] = np.linalg.norm(normed_preds - normed_targets)
            else:
                dists[c, n] = -1
    return dists


def dist_acc(dists, thr=0.5):
    '''
    Return percentage below threshold while ignoring values with a -1
    返回低于阈值的百分比，同时忽略带有 -1 的值
    '''
    dist_cal = np.not_equal(dists, -1)
    num_dist_cal = dist_cal.sum()
    if num_dist_cal > 0:
        return np.less(dists[dist_cal], thr).sum() * 1.0 / num_dist_cal
    else:
        return -1


def accuracy(output, target, hm_type='gaussian'):
    '''
    根据PCK计算精度，使用ground truth热力图而不是x,y坐标。
    第一个返回的值是所有关键点的平均精度，其次是单个精度。
    output:模型的预测输出
    target:样本的真实标签
    return:
    acc:比关键点数多一个，第0个存储这所有关键点的平均精度，后面的存储着单个关键点的精度
    avg_acc:,
    cnt:精度大于0的关键点数量
    pred:ndarray:(64,num_keypoints,2)
    '''
    num_keypoints = list(range(output.shape[1]))
    norm = 1.0
    if hm_type == 'gaussian':
        pred, _ = get_max_preds(output)#根据预测结果得到坐标信息
        target, _ = get_max_preds(target)#根据真实标签得到坐标信息
        h = output.shape[2]
        w = output.shape[3]
        norm = np.ones((pred.shape[0], 2)) * np.array([h, w]) / 10
    dists = calc_dists(pred, target, norm)#根据高斯核来判断差异

    acc = np.zeros((len(num_keypoints) + 1))
    avg_acc = 0
    cnt = 0

    for i in range(len(num_keypoints)):
        acc[i + 1] = dist_acc(dists[num_keypoints[i]])
        if acc[i + 1] >= 0:
            avg_acc += acc[i + 1]
            cnt += 1

    avg_acc = avg_acc / cnt if cnt != 0 else 0
    if cnt != 0:
        acc[0] = avg_acc
    return acc, avg_acc, cnt, pred
